import json
import os

from dotenv import load_dotenv
from fastapi import APIRouter
from pydantic import BaseModel

from ..models.ai_models.kg import build_knowledge_graph
kg=APIRouter(prefix="/kg", tags=["知识图谱"])

@kg.get("/")
async def get_video():
    return {"message": "Hello World"}

class KgParams(BaseModel):
    seed:str=...
    max_depth:int=...
    max_nodes:int=...

@kg.post("/get-kg")
async def get_kg(params:KgParams):
    """
    # 参数说明
    - seed:起始词
    - max_depth: 控制知识图谱扩展的最大深度
    - max_nodes: 限制知识图谱中最大节点数量
    """
    seed=params.seed
    max_depth=params.max_depth
    max_nodes=params.max_nodes

    graph=build_knowledge_graph(seed,max_depth,max_nodes)
    return {
        "code":200,
        "message":"成功",
        "data":graph
    }
